import copy
from typing import Union, Iterable, List, Tuple

import numpy as np

from ..trajectory_parameter import TrajectoryParameter
from ..trajectory_planner import TrajectoryPlanner


class TrajectoryBlendPlanner:
    def __init__(self, trajectory_parameters: Union[List[TrajectoryParameter], Tuple[TrajectoryParameter]],
                 radii: Union[np.ndarray, Iterable, int, float]) -> None:
        super().__init__()
        self.trajectory_parameters = copy.deepcopy(trajectory_parameters)
        self.radii = copy.deepcopy(radii)
        self.s_array = []
        self.tfs_cumsum = []
        self.planners: List[TrajectoryPlanner] = []
        self.calculate_radii()
        self.plan()

    def calculate_radii(self) -> None:
        for i, radius in enumerate(self.radii):
            max_radius = min((self.trajectory_parameters[i].get_tf(), self.trajectory_parameters[i + 1].get_tf())) / 2.0

            if radius > max_radius:
                self.radii[i] = max_radius

    def plan(self) -> None:
        tfs = [0.0, 0.0]
        for i in range(len(self.trajectory_parameters)):
            if i == 0:
                tf = self.trajectory_parameters[i].get_tf()
                if len(self.radii) != 0:
                    tf -= self.radii[i]
            elif i == len(self.trajectory_parameters) - 1:
                tf = self.trajectory_parameters[i].get_tf() - self.radii[i - 1]
            else:
                tf = self.trajectory_parameters[i].get_tf() - self.radii[i - 1] - self.radii[i]

            tfs.append(tf)

            if i != len(self.trajectory_parameters) - 1:
                tfs.append(self.radii[i])
        tfs.append(0.0)

        self.tfs_cumsum = np.cumsum(tfs)

        self.planners = [TrajectoryPlanner(trajectory_parameter) for trajectory_parameter in self.trajectory_parameters]

    def interpolate(self, t: float):

        if t <= 0.0:
            return self.planners[0].interpolate(0.0)
        elif t >= self.tfs_cumsum[-1]:
            return self.planners[-1].interpolate(self.trajectory_parameters[-1].get_tf())
        else:
            for i, ei in enumerate(self.tfs_cumsum[2:-1]):
                if t > ei:
                    continue

                if i % 2:
                    tt = (t - self.tfs_cumsum[i + 1]) / (
                            self.tfs_cumsum[i + 2] - self.tfs_cumsum[i + 1])
                    alpha = 6 * tt ** 5 - 15 * tt ** 4 + 10 * tt ** 3

                    t1 = (t - self.tfs_cumsum[i - 1])
                    t2 = (t - self.tfs_cumsum[i + 1])

                    return self.planners[i // 2].interpolate(t1) + alpha * (
                            self.planners[i // 2 + 1].interpolate(t2) - self.planners[i // 2].interpolate(t1))
                else:
                    tt = (t - self.tfs_cumsum[i])

                    return self.planners[i // 2].interpolate(tt)
